#include "PoissonSolver.h"
#include <cmath>
#include <iostream>

PoissonSolver::PoissonSolver(Grid grid, BoundaryConditions bc, std::function<double(double, double)> f)
    : grid(grid), bc(bc), f(f) {
    int n = grid.getN();
    solution = Eigen::VectorXd::Zero((n+2) * (n+2));
}

void PoissonSolver::solve() {
    // 获取网格步长和大小
    double h = grid.geth();
    int n = grid.getN();  

    // 点数
    int interiorPoints = (n + 2) * (n + 2);

    // 初始化系数矩阵 A 和右端项 b
    Eigen::SparseMatrix<double> A(interiorPoints, interiorPoints);
    Eigen::VectorXd b(interiorPoints);
    A.reserve(Eigen::VectorXi::Constant(interiorPoints, 5));  // 每行最多 5 个非零元素

    if (bc.getType() == "Dirichlet" && bc.hasHole() == false){
        // 填充系数矩阵 A 和右端项 b
        for (int j = 0; j < n + 2; ++j) {
            for (int i = 0; i < n + 2; ++i) {
                int index = j * (n + 2) + i;  // 将二维索引转换为一维索引

                if(i==0 || i==n+1 || j==0 || j==n+1){
                    // 边界时
                    A.insert(index, index) = 1.0;

                    // 边界条件
                    b(index) = bc.getCondition(grid.getX(i),grid.getY(j),0);
                }

                else{
                    // 中心点
                    A.insert(index, index) = 4.0 / (h * h);

                    // 下邻居
                    A.insert(index, index - (n+2)) = -1.0 / (h * h);

                    // 上邻居
                    A.insert(index, index + (n+2)) = -1.0 / (h * h);

                    // 左邻居
                    A.insert(index, index - 1) = -1.0 / (h * h);

                    // 右邻居
                    A.insert(index, index + 1) = -1.0 / (h * h);

                    // 右端项 b
                    b(index) = f(grid.getX(i),grid.getY(j));
                }
            }
        }

        // 压缩稀疏矩阵
        A.makeCompressed();
    }

    else if (bc.getType() == "Neumann" && bc.hasHole() == false) {
        // 填充系数矩阵 A 和右端项 b
        for (int j = 0; j < n + 2; ++j) {
            for (int i = 0; i < n + 2; ++i) {
                int index = j * (n + 2) + i;  // 将二维索引转换为一维索引

                // 四角上的点
                if (i == 0 && j == 0) {
                    // 左下角时
                    A.insert(index, index) = 1.0;
                    b(index) = std::exp(grid.getY(j) + std::sin(grid.getX(i)));
                }
                else if (i == n + 1 && j == 0) {
                    // 右下角时
                    A.insert(index, index) = 4.0 / (h * h);
                    A.insert(index, index + (n + 2)) = -2.0 / (h * h);  // 上邻居
                    A.insert(index, index - 1) = -2.0 / (h * h);  // 左邻居
                    b(index) = f(grid.getX(i), grid.getY(j)) + 2.0 * bc.getCondition(grid.getX(i), grid.getY(j), 1) / h + 2.0 * bc.getCondition(grid.getX(i), grid.getY(j), 2) / h;
                }
                else if (i == 0 && j == n + 1) {
                    // 左上角时
                    A.insert(index, index) = 4.0 / (h * h);
                    A.insert(index, index - (n + 2)) = -2.0 / (h * h);  // 下邻居
                    A.insert(index, index + 1) = -2.0 / (h * h);  // 右邻居
                    b(index) = f(grid.getX(i), grid.getY(j)) + 2.0 * bc.getCondition(grid.getX(i), grid.getY(j), 0) / h + 2.0 * bc.getCondition(grid.getX(i), grid.getY(j), 3) / h;
                }
                else if (i == n + 1 && j == n + 1) {
                    // 右上角时
                    A.insert(index, index) = 4.0 / (h * h);
                    A.insert(index, index - (n + 2)) = -2.0 / (h * h);  // 下邻居
                    A.insert(index, index - 1) = -2.0 / (h * h);  // 左邻居
                    b(index) = f(grid.getX(i), grid.getY(j)) + 2.0 * bc.getCondition(grid.getX(i), grid.getY(j), 1) / h + 2.0 * bc.getCondition(grid.getX(i), grid.getY(j), 3) / h;
                }

                // 四边上的点（不包括四个角）
                else if (j == 0 && i != 0 && i != n + 1) {
                    // 下边时
                    A.insert(index, index) = 4.0 / (h * h);
                    A.insert(index, index + (n + 2)) = -2.0 / (h * h);  // 上邻居
                    A.insert(index, index - 1) = -1.0 / (h * h);  // 左邻居
                    A.insert(index, index + 1) = -1.0 / (h * h);  // 右邻居
                    b(index) = f(grid.getX(i), grid.getY(j)) + 2.0 * bc.getCondition(grid.getX(i), grid.getY(j), 2) / h;
                }
                else if (i == 0 && j != 0 && j != n + 1) {
                    // 左边时
                    A.insert(index, index) = 4.0 / (h * h);
                    A.insert(index, index - (n + 2)) = -1.0 / (h * h);  // 下邻居
                    A.insert(index, index + (n + 2)) = -1.0 / (h * h);  // 上邻居
                    A.insert(index, index + 1) = -2.0 / (h * h);  // 右邻居
                    b(index) = f(grid.getX(i), grid.getY(j)) + 2.0 * bc.getCondition(grid.getX(i), grid.getY(j), 0) / h;
                }
                else if (i == n + 1 && j != 0 && j != n + 1) {
                    // 右边时
                    A.insert(index, index) = 4.0 / (h * h);
                    A.insert(index, index - (n + 2)) = -1.0 / (h * h);  // 下邻居
                    A.insert(index, index + (n + 2)) = -1.0 / (h * h);  // 上邻居
                    A.insert(index, index - 1) = -2.0 / (h * h);  // 左邻居
                    b(index) = f(grid.getX(i), grid.getY(j)) + 2.0 * bc.getCondition(grid.getX(i), grid.getY(j), 1) / h;
                }
                else if (j == n + 1 && i != 0 && i != n + 1) {
                    // 上边时
                    A.insert(index, index) = 4.0 / (h * h);
                    A.insert(index, index - (n + 2)) = -2.0 / (h * h);  // 下邻居
                    A.insert(index, index + 1) = -1.0 / (h * h);  // 右邻居
                    A.insert(index, index - 1) = -1.0 / (h * h);  // 左邻居
                    b(index) = f(grid.getX(i), grid.getY(j)) + 2.0 * bc.getCondition(grid.getX(i), grid.getY(j), 3) / h;
                }

                // 内部点
                else {
                    // 中心点
                    A.insert(index, index) = 4.0 / (h * h);
                    A.insert(index, index - (n + 2)) = -1.0 / (h * h);  // 下邻居
                    A.insert(index, index + (n + 2)) = -1.0 / (h * h);  // 上邻居
                    A.insert(index, index - 1) = -1.0 / (h * h);  // 左邻居
                    A.insert(index, index + 1) = -1.0 / (h * h);  // 右邻居
                    b(index) = f(grid.getX(i), grid.getY(j));
                }
            }
        }

        // 压缩稀疏矩阵
        A.makeCompressed();
    }


    // 打印矩阵 A 和向量 b（调试用）
    /*std::cout << "Matrix A:\n" << Eigen::MatrixXd(A) << std::endl;
    std::cout << "Vector b:\n" << b << std::endl;*/

    // 使用LU分解求解线性方程组 A * solution = b
    Eigen::SparseLU<Eigen::SparseMatrix<double>> solver;
    solver.compute(A);
    
    // 检查LU分解是否成功
    if (solver.info() != Eigen::Success) {
        std::cerr << "Failed to decompose the matrix!" << std::endl;
        return;
    }
    
    // 求解线性方程组
    solution = solver.solve(b);
    
    // 检查解是否成功
    if (solver.info() != Eigen::Success) {
        std::cerr << "Failed to solve the linear system!" << std::endl;
        return;
    }
    
    // 输出解
    /*std::cout << "Solution: \n" << solution << std::endl;*/
}

Eigen::VectorXd PoissonSolver::getSolution() const {
    return solution;
}